# -*- coding : utf-8 -*-


import io
import os
import json
import tqdm
import uuid
# import torch
import random
import transformers

from openpyxl import load_workbook, Workbook


def generate_uuid(): return str(uuid.uuid4())

def _make_w_io_base(f, mode: str):
    if not isinstance(f, io.IOBase):
        f_dirname = os.path.dirname(f)
        if f_dirname != "":
            os.makedirs(f_dirname, exist_ok=True)
        f = open(f, mode=mode)
    return f


def _make_r_io_base(f, mode: str):
    if not isinstance(f, io.IOBase):
        f = open(f, mode=mode)
    return f


def jdump(obj, f, mode="w", indent=4, default=str):
    """Dump a str or dictionary to a file in json format.

    Args:
        obj: An object to be written.
        f: A string path to the location on disk.
        mode: Mode for opening the file.
        indent: Indent for storing json dictionaries.
        default: A function to handle non-serializable entries; defaults to `str`.
    """
    f = _make_w_io_base(f, mode)
    if isinstance(obj, (dict, list)):
        json.dump(obj, f, indent=indent, default=default, ensure_ascii=False)
    elif isinstance(obj, str):
        f.write(obj)
    else:
        raise ValueError(f"Unexpected type: {type(obj)}")
    f.close()


def jload(f, mode="r"):
    """Load a .json file into a dictionary."""
    f = _make_r_io_base(f, mode)
    jdict = json.load(f)
    f.close()
    return jdict


def original_train_data():

    wb = load_workbook('./FAQ.xlsx')
    ws = wb[wb.sheetnames[0]]

    original_intent_dict = dict()
    other_intent_dict = dict()
    for i, row in enumerate(ws.values):
        if i != 0:
            intent_name = row[1]
            query = row[2]

            if intent_name == 'NOINTENT':
                intent_name = row[6]
            if intent_name in ['在操作', ]:
                # continue
                intent_name = '正在操作'
            if intent_name in ['在忙-无原因', '在忙-有原因', ]:
                intent_name = '没时间'
            if intent_name == '在忙-快点说':
                intent_name = '快点说'
            if intent_name == '在忙-主动邀约':
                intent_name = '主动邀约'

            if i != 0 and row[5] in [1, 2]:
                if intent_name not in original_intent_dict:
                    original_intent_dict[intent_name] = [list(), '']
                original_intent_dict[intent_name][0].append(query)
            else:
                if intent_name not in other_intent_dict:
                    other_intent_dict[intent_name] = [list(), '']
                other_intent_dict[intent_name][0].append(query)

    ws = wb[wb.sheetnames[1]]
    for i, row in enumerate(ws.values):
        intent_name = row[0]
        intent_meaning = row[1]
        # use_sign = 1 if row[5] == 1 else 0

        if intent_name in original_intent_dict:
            original_intent_dict[intent_name][1] = intent_meaning

    return original_intent_dict, other_intent_dict


def get_train_data():

    intent_dict, _ = original_train_data()
    data_list = list()

    instruction = '''接下来会给你提供一个用户的表述，你需要做的是用一个词代表用户的表述含义。'''
    for intent_name in intent_dict:
        for query in intent_dict[intent_name][0]:
            data_list.append({
                'instruction': instruction,
                'input': '用户：'+query,
                'output': intent_name
            })

    print(len(data_list))
    random.shuffle(data_list)
    jdump(data_list, './nlu_train_data.json')


def generate_prompt(instruction, input=None):
    if input:
        return f"""Instruction:\n{instruction}\n\n### Input:{input}\n\n### Response:"""
    else:
        return f"""Instruction:\n{instruction}\n\n### Response:"""


def get_test_data():
    model_path = './bloom-1.7b-nlu'
    tokenizer = transformers.AutoTokenizer.from_pretrained(model_path, padding_side="right")
    model = transformers.AutoModelForCausalLM.from_pretrained(model_path)
    model = model.to("cuda:0")

    num = 0
    # _, intent_dict = original_train_data() # 第一次迭代
    intent_dict = itera_model()
    wb = Workbook()
    ws_1 = wb.active
    ws_1.append([
        '语料', '意图名称',
    ])
    ws_2 = wb.create_sheet('sheet2')
    ws_2.append([
        '语料', '意图名称', '迭代意图名称'
    ])

    request_list = list()
    for intent_name in tqdm.tqdm(intent_dict):
        for data in tqdm.tqdm(intent_dict[intent_name][0]):
            instruction = '''接下来会给你提供一个用户的表述，你需要做的是用一个词代表用户的表述含义。'''
            prompt = generate_prompt(instruction, data)

            inputs = tokenizer(prompt, return_tensors="pt")
            input_ids = inputs["input_ids"].to("cuda:0")
            with torch.no_grad():
                generation_output = model.generate(
                    input_ids=input_ids,
                    temperature=0.1,
                    top_p=0.95,
                    top_k=1,
                    do_sample=False,
                    num_beams=1,
                    max_new_tokens=100,
                    eos_token_id=tokenizer.eos_token_id,
                    pad_token_id=tokenizer.pad_token_id,
                    return_dict_in_generate=True,
                    output_scores=True
                )

                s = generation_output.sequences[0]
                model_output = tokenizer.decode(s)

                res = model_output.split("### Response:")[-1].strip().replace('</s>', '')

                if intent_name == res:
                    ws_1.append([data, intent_name, ])
                else:
                    # print(model_output)
                    # print(intent_name)
                    ws_2.append([data, intent_name, res])

    wb.save('./第五次迭代.xlsx')


def itera_model():

    intent_dict, _ = original_train_data()
    test_intent_dict = dict()
    wb_first = load_workbook('./第一次迭代.xlsx')
    ws = wb_first[wb_first.sheetnames[0]]
    for i, row in enumerate(ws.values):
        if i != 0:
            query = row[0]
            intent = row[1]

            if intent not in intent_dict:
                intent_dict[intent] = [list(), '']
            if query not in intent_dict[intent][0]:
                intent_dict[intent][0].append(query)

    ws = wb_first[wb_first.sheetnames[1]]
    for i, row in enumerate(ws.values):
        if i != 0 and row[0]:
            query = row[0]
            intent = row[1]
            model_intent = row[2]

            if not intent and model_intent in ['不明', '肯定', '否定']:
                if model_intent not in intent_dict:
                    intent_dict[model_intent] = [list(), '']
                if query not in intent_dict[model_intent][0]:
                    intent_dict[model_intent][0].append(query)

            '''else:
                if intent not in test_intent_dict:
                    test_intent_dict[intent] = [list(), '']
                test_intent_dict[intent][0].append(query)'''

    wb_first = load_workbook('./第二次迭代.xlsx')
    ws = wb_first[wb_first.sheetnames[0]]
    for i, row in enumerate(ws.values):
        if i != 0:
            query = row[0]
            intent = row[1]

            if intent not in intent_dict:
                intent_dict[intent] = [list(), '']
            if query not in intent_dict[intent][0]:
                intent_dict[intent][0].append(query)

    ws = wb_first[wb_first.sheetnames[1]]
    for i, row in enumerate(ws.values):
        if i != 0 and row[0]:
            query = row[0]
            intent = row[1]
            model_intent = row[2]

            if not intent and model_intent in ['不明', '肯定', '否定']:
                if model_intent not in intent_dict:
                    intent_dict[model_intent] = [list(), '']
                if query not in intent_dict[model_intent][0]:
                    intent_dict[model_intent][0].append(query)

            '''else:
                if intent not in test_intent_dict:
                    test_intent_dict[intent] = [list(), '']
                test_intent_dict[intent][0].append(query)'''

    wb_first = load_workbook('./第三次迭代.xlsx')
    ws = wb_first[wb_first.sheetnames[0]]
    for i, row in enumerate(ws.values):
        if i != 0:
            query = row[0]
            intent = row[1]

            if intent not in intent_dict:
                intent_dict[intent] = [list(), '']
            if query not in intent_dict[intent][0]:
                intent_dict[intent][0].append(query)

    ws = wb_first[wb_first.sheetnames[1]]
    for i, row in enumerate(ws.values):
        if i != 0 and row[0]:
            query = row[0]
            intent = row[1]
            model_intent = row[2]

            if not intent and model_intent in ['不明', '肯定', '否定']:
                if model_intent not in intent_dict:
                    intent_dict[model_intent] = [list(), '']
                if query not in intent_dict[model_intent][0]:
                    intent_dict[model_intent][0].append(query)

            '''else:
                if intent not in test_intent_dict:
                    test_intent_dict[intent] = [list(), '']
                test_intent_dict[intent][0].append(query)'''

    wb_first = load_workbook('./第四次迭代.xlsx')
    ws = wb_first[wb_first.sheetnames[0]]
    for i, row in enumerate(ws.values):
        if i != 0:
            query = row[0]
            intent = row[1]

            if intent not in intent_dict:
                intent_dict[intent] = [list(), '']
            if query not in intent_dict[intent][0]:
                intent_dict[intent][0].append(query)

    ws = wb_first[wb_first.sheetnames[1]]
    for i, row in enumerate(ws.values):
        if i != 0 and row[0]:
            query = row[0]
            intent = row[1]
            model_intent = row[2]

            if not intent and model_intent in ['不明', '肯定', '否定']:
                if model_intent not in intent_dict:
                    intent_dict[model_intent] = [list(), '']
                if query not in intent_dict[model_intent][0]:
                    intent_dict[model_intent][0].append(query)

            '''else:
                if intent not in test_intent_dict:
                    test_intent_dict[intent] = [list(), '']
                test_intent_dict[intent][0].append(query)'''

    wb_first = load_workbook('./第五次迭代.xlsx')
    ws = wb_first[wb_first.sheetnames[0]]
    for i, row in enumerate(ws.values):
        if i != 0:
            query = row[0]
            intent = row[1]

            if intent not in intent_dict:
                intent_dict[intent] = [list(), '']
            if query not in intent_dict[intent][0]:
                intent_dict[intent][0].append(query)

    ws = wb_first[wb_first.sheetnames[1]]
    for i, row in enumerate(ws.values):
        if i != 0 and row[0]:
            query = row[0]
            intent = row[1]
            model_intent = row[2]

            if not intent and model_intent in ['不明', '肯定', '否定']:
                if model_intent not in intent_dict:
                    intent_dict[model_intent] = [list(), '']
                if query not in intent_dict[model_intent][0]:
                    intent_dict[model_intent][0].append(query)

            else:
                if intent not in test_intent_dict:
                    test_intent_dict[intent] = [list(), '']
                test_intent_dict[intent][0].append(query)

    data_list = list()
    instruction = '''接下来会给你提供一个用户的表述，你需要做的是用一个词代表用户的表述含义。'''
    for intent_name in intent_dict:
        for query in intent_dict[intent_name][0]:
            '''data_list.append({
                'instruction': instruction,
                'input': '用户：'+query,
                'output': intent_name
            })'''

            data_list.append({
            'id': generate_uuid(),
            'model': '',
            'conversations': [{
                'from': 'human',
                'value': query,
            },
            {
                'from': 'gpt',
                'value': intent_name,
            }]
        })

    print(len(data_list))
    random.shuffle(data_list)
    jdump(data_list, './nlu_train_data.json')

    return test_intent_dict


def test_data():
    model_path = './bloom-1.7b-nlu'
    tokenizer = transformers.AutoTokenizer.from_pretrained(model_path, padding_side="right")
    model = transformers.AutoModelForCausalLM.from_pretrained(model_path)
    model = model.to("cuda:0")

    wb = load_workbook('./标注测试.xlsx')
    ws = wb[wb.sheetnames[0]]

    wb_w = Workbook()
    ws_w = wb_w.active
    ws_w.append([
        '语句', '线上意图', '模型意图'
    ])

    attitude_dict = {
        -1: '否定',
        0: '不明',
        1: '肯定',
    }
    for i, row in enumerate(ws.values):
        if i != 0:
            query = row[2]
            intent = row[3]
            attitude = row[4]

            if intent == 'NOINTENT':
                intent = attitude_dict[attitude]
            elif intent == '肯定态度':
                intent = '肯定'
            elif intent == '否定态度':
                intent = '否定'

            instruction = '''接下来会给你提供一个用户的表述，你需要做的是用一个词代表用户的表述含义。'''
            prompt = generate_prompt(instruction, '用户：'+query)

            inputs = tokenizer(prompt, return_tensors="pt")
            input_ids = inputs["input_ids"].to("cuda:0")
            with torch.no_grad():
                generation_output = model.generate(
                    input_ids=input_ids,
                    temperature=0.1,
                    top_p=0.95,
                    top_k=1,
                    do_sample=False,
                    num_beams=1,
                    max_new_tokens=100,
                    eos_token_id=tokenizer.eos_token_id,
                    pad_token_id=tokenizer.pad_token_id,
                    return_dict_in_generate=True,
                    output_scores=True
                )

                s = generation_output.sequences[0]
                model_output = tokenizer.decode(s)

                res = model_output.split("### Response:")[-1].strip().replace('</s>', '')
                ws_w.append([
                    query, intent, res
                ])

    wb_w.save('./迭代测试结果.xlsx')


if __name__ == '__main__':


    # get_test_data()

    itera_model()

    # test_data()
